import random
from typing import Dict, Union
import os
import json
import ollama
from pydantic import BaseModel, Field

# Ollama配置
client = ollama.Client(host='http://localhost:11434')


class WeatherInfo(BaseModel):
    temperature: float = Field(..., description="当前温度（摄氏度）")
    feels_like: float = Field(..., description="体感温度（摄氏度）")
    description: str = Field(..., description="天气描述")
    humidity: int = Field(..., description="湿度百分比")
    wind_speed: float = Field(..., description="风速(m/s)")


class ChatService:
    def __init__(self):
        self.model = "llama3.2:latest"

    def process_message(self, message: str) -> Dict:
        """处理用户消息，识别意图"""
        system_prompt = """你是一个智能助手，主要负责理解用户的意图。
        如果用户询问天气相关信息，请提取出城市名称，并返回格式为: {"type": "weather", "city": "城市名"}
        如果是普通对话，请返回格式为: {"type": "chat", "message": "用户原始消息"}
        """

        response = client.chat(model=self.model, messages=[
            {
                'role': 'system',
                'content': system_prompt
            },
            {
                'role': 'user',
                'content': message
            }
        ])

        try:
            return json.loads(response['message']['content'])
        except:
            return {"type": "chat", "message": message}

    def chat_response(self, message: str) -> str:
        """处理普通对话"""
        system_prompt = "你是一个友好的助手，请用自然、亲切的方式回答用户的问题。"
        response = client.chat(model=self.model, messages=[
            {
                'role': 'system',
                'content': system_prompt
            },
            {
                'role': 'user',
                'content': message
            }
        ])
        return response['message']['content']


def format_weather_data(weather: WeatherInfo) -> Dict:
    """将天气信息格式化为显示格式"""
    return {
        "温度": f"{weather.temperature}°C",
        "体感温度": f"{weather.feels_like}°C",
        "天气": weather.description,
        "湿度": f"{weather.humidity}%",
        "风速": f"{weather.wind_speed}m/s"
    }


class WeatherService:
    def __init__(self):
        self.model = "llama3.2:latest"
        self._setup_fake_data()

    def _setup_fake_data(self):
        """设置假数据"""
        self.weather_descriptions = [
            "晴朗", "多云", "阴天", "小雨", "中雨",
            "大雨", "雷阵雨", "小雪", "中雪", "大雪"
        ]

    def get_weather(self, city: str) -> Union[WeatherInfo, str]:
        """获取模拟的天气信息"""
        try:
            # 生成随机天气数据
            temperature = round(random.uniform(15, 30), 1)  # 15-30度之间的随机温度
            feels_like = temperature + random.uniform(-2, 2)  # 体感温度与实际温度相差±2度
            description = random.choice(self.weather_descriptions)  # 随机天气描述
            humidity = random.randint(30, 90)  # 30-90%之间的随机湿度
            wind_speed = round(random.uniform(0, 10), 1)  # 0-10m/s之间的随机风速

            return WeatherInfo(
                temperature=temperature,
                feels_like=feels_like,
                description=description,
                humidity=humidity,
                wind_speed=wind_speed
            )
        except Exception as e:
            return f"获取天气信息失败: {str(e)}"

    def describe_weather(self, city: str, weather_data: Dict) -> str:
        """使用LLM生成天气描述"""
        system_prompt = "你是一个天气助手，请根据提供的天气数据，用自然的语言描述天气情况。"
        message = f"城市：{city}\n天气数据：{weather_data}"

        response = client.chat(model=self.model, messages=[
            {
                'role': 'system',
                'content': system_prompt
            },
            {
                'role': 'user',
                'content': message
            }
        ])
        return response['message']['content']


def chat():
    weather_service = WeatherService()
    chat_service = ChatService()

    print("你好！我是智能助手，可以帮你查询天气，也可以和你聊天。输入 'q' 退出。")

    while True:
        user_input = input("\n你: ")
        if user_input.lower() == 'q':
            print("再见！")
            break

        # 处理用户输入，识别意图
        intent = chat_service.process_message(user_input)

        if intent["type"] == "weather":
            # 处理天气查询
            weather_result = weather_service.get_weather(intent["city"])
            if isinstance(weather_result, WeatherInfo):
                formatted_data = format_weather_data(weather_result)
                response = weather_service.describe_weather(intent["city"], formatted_data)
            else:
                response = weather_result
        else:
            # 处理普通对话
            response = chat_service.chat_response(intent["message"])

        print(f"\n助手: {response}")


if __name__ == '__main__':
    chat()
